#!/bin/bash
export PYTHONUNBUFFERED=1

model=qwen1.8b.chat
new_date=$(date +%Y-%m-%d)_${model}
if [ ! -d "log/$new_date" ]; then
    mkdir -p "log/$new_date"
fi

model_name=Qwen-1_8B-Chat
data_type=CKnowEdit
# data_path=../../EasyEdit/dataset/ccks2024_know_edit/ccks-CKnowEdit.json

source activate ke2torch23cu121
export HUGGINGFACE_CACHE=/share/huggingface/

batch_size=10
prob_type=target_new
layer_idx=0,24
max_tokens=128
max_sample_steps=320

i=1
ff_attrs=attn.c_attn
for data_name in {type1_133,type2_80,type4_50,type5_70,type6_50}; do
    echo "$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps"
    CUDA_VISIBLE_DEVICES=$i python kn.py \
        --model_name $model_name \
        --batch_size $batch_size \
        --data_type $data_type \
        --data_path ../../EasyEdit/dataset/ccks2024_know_edit/CKnowEditType/$data_name.json \
        --transformer_layers_attr transformer.h \
        --ff_attrs $ff_attrs \
        --prob_type $prob_type \
        --layer_idx $layer_idx \
        --steps 20 \
        --max_tokens $max_tokens \
        --max_sample_steps $max_sample_steps \
        > log/$new_date/$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps-1.log 2>&1 &
    wait
done

# i=0
# data_name=type3_40
# for ff_attrs in {mlp.c_proj,attn.c_proj,mlp.w1,mlp.w2,attn.c_attn,}; do
#     echo "$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps"
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path /share/dataset/CKnowEditType/$data_name.json \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps 20 \
#         --max_tokens $max_tokens \
#         --max_sample_steps $max_sample_steps \
#         > log/$new_date/$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps-1.log 2>&1 &
#     wait
# done


# i=1
# ff_attrs=mlp.c_proj
# for data_name in {type2_80,type3_40,type4_50,type5_70,type6_50}; do
#     echo "$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps"
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path ../../EasyEdit/dataset/ccks2024_know_edit/CKnowEditType/$data_name.json \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps 20 \
#         --max_tokens $max_tokens \
#         --max_sample_steps $max_sample_steps \
#         > log/$new_date/$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps-1.log 2>&1 &
#     wait
# done

# i=2
# ff_attrs=attn.c_proj
# for data_name in {type2_80,type3_40,type4_50,type5_70,type6_50}; do
#     echo "$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps"
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path ../../EasyEdit/dataset/ccks2024_know_edit/CKnowEditType/$data_name.json \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps 20 \
#         --max_tokens $max_tokens \
#         --max_sample_steps $max_sample_steps \
#         > log/$new_date/$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps-1.log 2>&1 &
#     wait
# done

# i=3
# ff_attrs=attn.c_attn
# for data_name in {type2_80,type3_40,type4_50,type5_70,type6_50}; do
#     echo "$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps"
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path ../../EasyEdit/dataset/ccks2024_know_edit/CKnowEditType/$data_name.json \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps 20 \
#         --max_tokens $max_tokens \
#         --max_sample_steps $max_sample_steps \
#         > log/$new_date/$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps-1.log 2>&1 &
#     wait
# done


# skip=3
# for i in $(seq 0 7); do
#     idx1=$(($i * $skip))
#     idx2=$(($i * $skip + $skip))
#     layer_idx=$idx1,$idx2
#     echo "$i-$model-$data_type-$layer_idx-$start_idx-$end_idx-bs$batch_size-$prob_type-$ff_attrs"
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --start_idx $start_idx \
#         --end_idx  $end_idx \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path $data_path \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps 20 \
#         > log/$new_date/$i-$model-$data_type-$layer_idx-$start_idx-$end_idx-bs$batch_size-$prob_type-$ff_attrs-7.log 2>&1 &
# done


# i=2
# for ff_attrs in {mlp.w1,mlp.w2,attn.c_attn,attn.c_proj}; do
#     echo "$i $model_name $data_type $batch_size $prob_type $ff_attrs $layer_idx"
#     start_idx=423
#     end_idx=700 # 700-277
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --start_idx $start_idx \
#         --end_idx  $end_idx \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path $data_path \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps $batch_size \
#         --next_token answer_next_token \
#         > log/$new_date/$i-$model-$data_type-$layer_idx-$start_idx-$end_idx-bs$batch_size-$prob_type-$ff_attrs-1.log 2>&1 &
#     wait
# done


# for ff_attrs in {mlp.w1,mlp.w2,attn.c_attn,attn.c_proj}; do
#     echo "$i $model_name $data_type $batch_size $prob_type $ff_attrs $layer_idx $start_idx $end_idx"
#     CUDA_VISIBLE_DEVICES=$i python kn.py \
#         --model_name $model_name \
#         --start_idx $start_idx \
#         --end_idx  $end_idx \
#         --batch_size $batch_size \
#         --data_type $data_type \
#         --data_path $data_path \
#         --transformer_layers_attr transformer.h \
#         --ff_attrs $ff_attrs \
#         --prob_type $prob_type \
#         --layer_idx $layer_idx \
#         --steps $batch_size \
#         --next_token answer_next_token \
#         > log/$new_date/$i-$model-$data_type-$layer_idx-$start_idx-$end_idx-bs$batch_size-$prob_type-$ff_attrs-1.log 2>&1 &
#     wait
# done
